package cn.dataling.ai.component;

import cn.dataling.ai.enums.ChatModelEnum;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.util.Map;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

@Component
public class ChatComponent {

    @Resource
    private Map<String, ChatClient> chatClientMap;
    @Value("${currentModel}")
    private String currentModel;

    public Flux<String> chatByStream(String message,String msgId) {

        String serviceName = ChatModelEnum.getChatClientName(currentModel);
        ChatClient chatClient = chatClientMap.get(serviceName);
        return chatClient.prompt()
                .user(message)
                .advisors(a -> a
                        .param(CHAT_MEMORY_CONVERSATION_ID_KEY, msgId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 5)
                )
                .stream()
                .content();
    }
}
